#!/usr/bin/env python
#-*- coding:utf-8 -*-
import os,sys
import gzip
import bz2
import tarfile
import zipfile
import zlib
import urllib2
import glob
import re
from wwpy.util.Module import pickle,StringIO,xlrd,xlwt
from wwpy.util.Logger import logger
from wwpy.util.Threads import ThreadPool


class BaseSource(object):
    
    def __init__(self,source):
        self.source=source
        self._sep=None
        self._line_sep='\n'
        self._info=self._init_info()
        
    def _init_info(self):
        source=self.source
        return dict(
            source=source,
            handler=None,
            member=source,
            members=[source],
            source_type=None,
            coding='utf-8',
        )
    
    @property
    def info(self):
        return self._info
    
    @property
    def is_compressed(self):
        return self.source_type in '.GZ,.BZ2,.ZIP,.TAR,.TAR.GZ,.TAR.BZ2'.split(',')
    
    @property
    def is_tar(self):
        return self.source_type.startswith('.TAR')
    
    @property
    def is_zip(self):
        return self.source_type=='.ZIP'
    
    @property
    def is_file(self):
        if self.is_compressed:
            return True
        return self.source_type=='.FILE'
    
    @property
    def is_url(self):
        return self.source_type=='.URL'
    
    @property
    def is_sql(self):
        return self.source_type=='.SQL'
    
    @property
    def is_excel(self):
        return self.source_type=='.EXCEL'
    
    @property
    def is_std(self):
        return self.source_type=='.STD'
    
    @property
    def is_str(self):
        return self.source_type=='.STR'
    
    @property
    def handler(self):
        return self._info['handler']
    
    @property
    def members(self):
        return self._info['members']
    
    @property
    def member(self):
        return self._info['member']
    
    @property
    def source_type(self):
        return str(self._info['source_type'])
    
    @property
    def coding(self):
        return self._info['coding']
    
    def get_member(self,members,member=0):
        default_member=members[0]
        if isinstance(member,int):
            if 0<=member<len(members):
                return members[member]
            return default_member
        elif isinstance(member,basestring):
            ms=[m for m in members if m.lower()==member.lower()]
            if len(ms)>0:
                return ms[0]
            return default_member
        return default_member
    
    def default_line_func(self,line):
        if isinstance(line,basestring):
            if line.endswith(self._line_sep):
                line=line[:-len(self._line_sep)]
            if self._sep is not None:
                line=line.split(self._sep)
        return line
    
    def add_suffix_to_line(self,line):
        line=str(line)
        if not line.endswith(self._line_sep):
            line+=self._line_sep
        return line
    
    def to_string(self,line,sep='\t'):
        if isinstance(line,basestring):
            return line
        elif isinstance(line,(tuple,list,set)):
            return sep.join([(unicode(field) if isinstance(field,unicode) else str(field)) for field in line])
        return line
    
    def get_size(self):
        try:
            if self.is_file:
                return os.path.getsize(self.source)
            elif self.is_sql:
                return sum(sum(len(str(field)) for field in line) for line in self.open(mode='r'))
            elif any([self.is_url,self.is_std]):
                return sum(len(line) for line in self.open(mode='r'))
            elif self.is_str:
                return len(self.source.value)
        except:
            return 0
        
    def get_line_count(self):
        try:
            return sum(1 for line in self.open(mode='r'))
        except:
            return 0
    
    def open(self,**kwargs):
        pass
    
    def close(self):
        if hasattr(self.handler,'close'):
            self.handler.close()
        self._info=self._init_info()
        
    def read(self,**kwargs):
        mode=kwargs.pop('mode','r')
        debug=kwargs.pop('debug',True)
        line_func=kwargs.pop('line_func',None)
        try:
            self.open(mode=mode,**kwargs)
            if line_func is None:
                line_func=self.default_line_func
            if debug:
                if len(self.members)>1:
                    logger.info("Read source from [%s]->%s" % (self.source,self.member))
                else:
                    logger.info("Read source from [%s]" % self.source)
            for line in self.handler:
                line=line_func(line)
                if line is not None:
                    yield line
        except IOError,error:
            raise IOError(error)
        except Exception,error:
            raise Exception(error)
        
    def __iter__(self):
        return iter(self.read())
    
    def write(self,data,**kwargs):
        mode=kwargs.pop('mode','w')
        if self.handler is None:
            self.open(mode=mode,**kwargs)
        if self.is_compressed:
            return self.compress(data)
        else:
            line=self.add_suffix_to_line(data)
            self.handler.write(line)
        return self.handler
    
class FileSource(BaseSource):
    
    def __init__(self,source):
        super(FileSource,self).__init__(source)
        
    def open(self,**kwargs):
        f=open(self.source,**kwargs)
        self._info['handler']=f
        return f
    
    @property
    def source_type(self):
        self._info['source_type']='.FILE'
        return self._info['source_type']
    
class GzipSource(BaseSource):
    
    def __init__(self,source):
        super(GzipSource,self).__init__(source)
        
    def open(self,**kwargs):
        gz=gzip.open(self.source,**kwargs)
        self._info['handler']=gz
        return gz
    
    @property
    def source_type(self):
        self._info['source_type']='.GZ'
        return self._info['source_type']
    
    def compress(self,src,**kwargs):
        logger.info("Compress from %s into %s" % (src,self.source))
        s=Source(src)
        f=s.open(mode='r',**kwargs)
        for line in f:
            line=self.add_suffix_to_line(line)
            self.handler.write(line)
        s.close()
        
    def decompress(self,dest=None):
        if dest is None:
            dest=source[:-len(self.source_type)]
        logger.info('Decompress %s into %s' % (self.source,dest))
        f=open(dest,'w')
        for line in self.open(mode='r'):
            line=self.add_suffix_to_line(line)
            f.write(line)
        f.close()
        
class BZ2Source(BaseSource):
    
    def __init__(self,source):
        super(BZ2Source,self).__init__(source)
        
    def open(self,**kwargs):
        bz=bz2.BZ2File(self.source,**kwargs)
        self._info['handler']=bz
        return bz
    
    @property
    def source_type(self):
        self._info['source_type']='.BZ2'
        return self._info['source_type']
    
    def compress(self,src,**kwargs):
        logger.info("Compress from %s into %s" % (src,self.source))
        s=Source(src)
        f=s.open(mode='r',**kwargs)
        for line in f:
            line=self.add_suffix_to_line(line)
            self.handler.write(line)
        s.close()
        
    def decompress(self,dest=None):
        if dest is None:
            dest=source[:-len(self.source_type)]
        logger.info('Decompress %s into %s' % (self.source,dest))
        f=open(dest,'w')
        for line in self.open(mode='r'):
            line=self.add_suffix_to_line(line)
            f.write(line)
        f.close()
    
class ZipSource(BaseSource):
    
    def __init__(self,source):
        super(ZipSource,self).__init__(source)
        
    def open(self,**kwargs):
        member=kwargs.pop('member',0)
        mode=str(kwargs.get('mode','r')).lower()
        zf=zipfile.ZipFile(self.source,**kwargs)
        if mode.startswith('r'):
            members=zf.namelist()
            member=self.get_member(members,member)
            f=zf.open(member,mode,pwd=kwargs.get('pwd'))
            self._info['members']=members
            self._info['handler']=f
            self._info['member']=member
            return f
        elif mode.startswith('w') or mode.startswith('a'):
            self._info['handler']=zf
            return zf
        
    def compress(self,src,**kwargs):
        logger.info("Compress from %s into %s" % (src,self.source))
        self.handler.write(src)
    
    def decompress(self,path=None,**kwargs):
        if path is None:
            path=os.path.dirname(self.source)
        if not os.path.exists(path):
           os.mkdir(path)
        zf=zipfile.ZipFile(self.source,**kwargs)
        members=[]
        for name in zf.namelist():
            member=os.path.join(path,name)
            if member.endswith(os.sep):
                member=member.rstrip(os.sep)
                continue
            members.append(member)
            logger.info('Decompress %s->%s into %s' % (self.source,member,path))
            f=open(member,'wb')
            f.write(zf.read(name,pwd=path))
            f.close()
        zf.close()
        return members
        
    @property
    def source_type(self):
        self._info['source_type']='.ZIP'
        return self._info['source_type']
    
class TarSource(BaseSource):
    
    def __init__(self,source):
        super(TarSource,self).__init__(source)
    
    def open(self,**kwargs):
        mode=kwargs.get('mode','r')
        if 'mode' in kwargs: del kwargs['mode']
        if mode.startswith('r'):
            if self.source.endswith('.tar.gz'):
                mode='r:gz'
            elif self.source.endswith('.tar.bz2'):
                mode='r:bz2'
        elif mode.startswith('w'):
            if self.source.endswith('.tar.gz'):
                mode='w:gz'
            elif self.source.endswith('.tar.bz2'):
                mode='w:bz2'
        member=kwargs.pop('member',0)
        tar=tarfile.open(self.source,mode=mode,**kwargs)
        if mode.startswith('r'):
            members=tar.getnames()
            member=self.get_member(members,member)
            f=tar.extractfile(member)
            self._info['members']=members
            self._info['member']=member
            self._info['handler']=f
            return f
        elif mode.startswith('w'):
            self._info['handler']=tar
            return tar
        
    def compress(self,src):
        logger.info("Compress from %s into %s" % (src,self.source))
        self.handler.add(src)
        
    def decompress(self,path=None,**kwargs):
        if path is None:
            path=os.path.dirname(self.source)
        tar=tarfile.open(self.source,**kwargs)
        members=tar.getnames()
        for member in members:
            logger.info('Decompress %s->%s into %s' % (self.source,member,path))
            tar.extract(member,path=path)
        tar.close()
        return members
    
    @property
    def source_type(self):
        st='.TAR'
        if self.source.endswith('.tar.gz'):
            st+='.GZ'
        elif self.source.endswith('.tar.bz2'):
            st+='.BZ2'
        self._info['source_type']=st
        return self._info['source_type']
    
class ExcelSource(BaseSource):
    
    def __init__(self,source):
        super(ExcelSource,self).__init__(source)
        
    def default_line_func(self,line):
        return [unicode(field) for field in line]
        
    def get_member(self,members,member):
        default_member=members[0]
        if isinstance(member,int):
            if 0<=member<len(members):
                return members[member]
            return default_member
        elif isinstance(member,basestring):
            ms=[m for m in members if m.name.lower()==member.name.lower()]
            if len(ms)>0:
                return ms[0]
            return default_member
        return default_member
        
    def open(self,**kwargs):
        mode=str(kwargs.get('mode','r')).lower()
        member=kwargs.pop('member',0)
        if mode.startswith('r'):
            excel=xlrd.open_workbook(self.source)
            sheets=excel.sheets()
            sheet=self.get_member(sheets,member)
            rows=[]
            for i in range(sheet.nrows):
                rows.append(sheet.row_values(i))
            self._info['handler']=rows
            self._info['members']=[sheet.name for sheet in sheets]
            self._info['coding']='unicode'
            return rows
        elif mode.startswith('w'):
            excel=xlwt.Workbook()
            self._info['handler']=excel
            return excel
        
    @property
    def source_type(self):
        self._info['source_type']='.EXCEL'
        return self._info['source_type']
    
    def write(self,data,**kwargs):
        sep=self._sep or kwargs.pop('sep','\t')
        coding=kwargs.pop('coding','utf-8')
        self._sep=sep
        if self.handler is None:
            self.open(mode='w',**kwargs)
        s=Source(data)
        if s.is_sql or s.is_url:
            sheet_name='%s_%s' % (s.source_type[1:].lower(),str(abs(hash(data))))
        else:
            sheet_name=data
        excel=self.handler
        sheet=excel.add_sheet(sheet_name)
        def _decode_field(field):
            if isinstance(field,unicode):
                return field
            return field.decode(coding)
        for row,line in enumerate(s):
            if isinstance(line,basestring):
                fields=line.split(sep)
            else:
                fields=line
            for col,field in enumerate(fields):
                sheet.write(row,col,_decode_field(field))
        excel.save(self.source)
        
class StdStreamSource(BaseSource):
    
    def __init__(self,source=None):
        super(StdStreamSource,self).__init__(source)
        
    def open(self,**kwargs):
        mode=str(kwargs.get('mode','r')).lower()
        if mode.startswith('r'):
            f=sys.stdin
        elif mode.startswith('w'):
            f=sys.stdout
        self._info['handler']=f
        return f
    
    @property
    def source_type(self):
        self._info['source_type']='.STD'
        return self._info['source_type']
    
class UrlSource(BaseSource):
    
    def __init__(self,source):
        super(UrlSource,self).__init__(source)
        
    def open(self,**kwargs):
        mode=kwargs.pop('mode')
        f=urllib2.urlopen(self.source,**kwargs)
        self._info['handler']=f
        return f
    
    @property
    def source_type(self):
        self._info['source_type']='.URL'
        return self._info['source_type']
    
    def write(self,*args,**kwargs):
        pass
    
class SqlSource(BaseSource):
    def __init__(self,source):
        super(SqlSource,self).__init__(source)
        
    @property
    def dbs(self):
        from wwpy.db.DBC import dbc
        return dbc.keys()
    
    def default_line_func(self,line):
        return [(str(field) if field is not None else '') for field in line]
    
    def get_info(self):
        source=self.source
        regex=re.compile('^select\s+.+?\s+from\s+([^\s]+)',re.I|re.MULTILINE)
        m=regex.match(source)
        dbs=self.dbs
        db=None
        info={}
        if m:
            table=m.group(1)
            if table.find('.')>-1:
                _db,_table=table.split('.')
                if _db in dbs:
                    db=_db
                    source=re.sub(table,_table,source)
                    table=_table
        else:
            table=None
        info['source']=source
        info['db']=db
        info['table']=table
        return info
    
    def get_fields(self):
        info=self.get_info()
        db=info['db']
        if db is None:
            raise Exception("DBC must be in SQL,just like SELECT * FROM dbc.table!")
        sql=info['source']
        from wwpy.db.DBI import DBI,dbc
        dbi=DBI(**dbc[db])
        fields=dbi.fields(sql)
        dbi.close()
        return fields
    
    def open(self,**kwargs):
        from wwpy.db.DBI import DBI,dbc
        mode=str(kwargs.pop('mode','r'))
        info=self.get_info()
        source=info['source']
        db=info['db'] or (kwargs and kwargs.pop('db')) or 'db_base_dict'
        if mode.startswith('r'):
            dbi=DBI(**dbc[db])
            rows=dbi.fetch(source,**kwargs)
            self._info['db']=db
            self._info['table']=info['table']
            if kwargs.get('has_fields'):
                self._info['fields']=rows[0]
            else:
                self._info['fields']=dbi.fields()
            self._info['handler']=rows
            self._info['coding']=dbi.coding
            dbi.close()
            return rows
        else:
            pass
        
    def write(self,*args,**kwargs):
        pass
        
    @property
    def source_type(self):
        self._info['source_type']='.SQL'
        return self._info['source_type']
    
class StrSource(BaseSource):
    def __init__(self,source=None):
        super(StrSource,self).__init__(source)
        self.buffer=None
        
    def open(self,**kwargs):
        mode=str(kwargs.get('mode','r')).lower()
        if mode.startswith('r'):
            if self.source is not None:
                self.buffer=StringIO.StringIO(self.source)
            else:
                raise Exception('No string data supplied')
        elif mode.startswith('w'):
            if self.buffer is not None:
                self.buffer.close()
            self.buffer=StringIO.StringIO()
        elif mode.startswith('a'):
            if self.buffer is None:
                self.buffer=StringIO.StringIO()
        self._info['handler']=self.buffer
        return self.buffer
        
    @property
    def value(self):
        if self.buffer:
            return self.buffer.getvalue()
        return None
    
    @property
    def source_type(self):
        self._info['source_type']='.STR'
        return self._info['source_type']
        
class Source(object):
    
    def __init__(self,source=None):
        self._source=source
        self._sources=self._get_sources(source)
        self.source_objs=self._get_objs()
        self._set_attrs()
        
    def sep(self,sep=None):
        for obj in self.objs:
            obj._sep=sep
        return sep
    
    def line_sep(self,line_sep=None):
        for obj in self.objs:
            obj._line_sep=line_sep
        return line_sep
    
    def get(self):
        return self._source
    
    def set(self,source):
        self._source=source
        self._sources=self._get_sources(source)
        self.source_objs=self._get_objs()
        self._set_attrs()
        
    @property
    def is_dir(self):
        if isinstance(self._source,basestring):
            return os.path.isdir(self._source)
        return False
    
    def _get_sources(self,source):
        sources=[]
        if isinstance(source,basestring):
            if os.path.isdir(source):
                for src in os.listdir(source):
                    if src=='.' or src=='..':
                        continue
                    sources.append(os.path.join(source,src))
            else:
                ss=glob.glob(source)
                if not len(ss):
                    sources.append(source)
                else:
                    sources.extend(ss)
        elif isinstance(source,(tuple,list)):
            for src in source:
                sources.append(src)
        else:
            sources.append(source)
        sources.sort()
        return sources
        
    def _get_obj(self,source):
        if isinstance(source,basestring):
            if any(map(source.lower().startswith,['http://','https://','ftp://'])):
                return UrlSource(source)
            elif source.endswith('.tar.gz'):
                return TarSource(source)
            elif source.endswith('.tar.bz2'):
                return TarSource(source)
            elif source.endswith('.tar'):
                return TarSource(source)
            elif source.endswith('.gz'):
                return GzipSource(source)
            elif source.endswith('.bz2'):
                return BZ2Source(source)
            elif source.endswith('.zip'):
                return ZipSource(source)
            elif re.match(r'^select\s',source.strip(),re.I|re.MULTILINE):
                return SqlSource(source)
            elif any([source.endswith('.xls'),source.endswith('.xlsx')]):
                return ExcelSource(source)
            else:
                return FileSource(source)
        elif isinstance(source,BaseSource):
            return source
        elif source is None:
            return StdStreamSource()
        else:
            return source
        
    def _get_objs(self):
        objs={}
        for source in self._sources:
            obj=self._get_obj(source)
            objs[source]=obj
        return objs
    
    def _set_attrs(self):
        for obj in self.objs:
            methods=[method for method in dir(obj) if not method.startswith('_')]
            for method in methods:
                if not hasattr(self,method):
                    setattr(self,method,getattr(obj,method))
                   
    @property 
    def objs(self):
        return list(set(self.source_objs.values()))
    
    @property
    def sources(self):
        return self._sources
    
    @property
    def source(self):
        return self._sources[0]
    
    @property
    def obj(self):
        return self.source_objs[self.source]
    
    def close(self):
        for obj in self.objs:
            obj.close()
            
    def read(self,**kwargs):
        threads=kwargs.pop('threads',None)
        if len(self.sources)<2:
            threads=None
        if threads is not None:
            do_read=lambda obj:[line for line in obj.read(**kwargs)]
            try:
                tn=int(threads)
            except:
                tn=4
            tp=ThreadPool(tn)
            for obj in self.objs:
                tp.put(do_read,obj)
            tp.wait()
            for out in tp.outs():
                for line in out[-1]:
                    yield line
        else:
            for obj in self.objs:
                for line in obj.read(**kwargs):
                    yield line
                    
    def __iter__(self):
        for obj in self.objs:
            for line in obj:
                yield line
    
    def scan(self,**kwargs):
        sep=kwargs.pop('sep','\t')
        for line in self.read(**kwargs):
            print self.obj.to_string(line,sep=sep)
    
    def write(self,data,**kwargs):
        return self.obj.write(data,**kwargs)
    
    def compress(self,dest=None):
        if dest is None:
            dest="%s.gz" % self.source
        s=Source(dest)
        s.open(mode='w')
        for source in self.sources:
            s.obj.compress(source)
        s.close()
        return dest
    
    def decompress(self,dest=None):
        if not self.obj.is_compressed:
            logger.info("File %s is compressed." % (self.source))
            return
        return self.obj.decompress(dest)
    
